package edu.zju.gis.dbfg.model.statistic;

import edu.zju.gis.dbfg.common.exception.ModelFailedException;
import edu.zju.gis.dbfg.model.PyramidConfig;
import edu.zju.gis.dbfg.model.datasource.DataReaderRDD;
import edu.zju.gis.dbfg.model.util.CrsPick;
import edu.zju.gis.dbfg.model.util.GeoItem;
import edu.zju.gis.dbfg.model.util.TileID;
import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.api.java.function.Function;
import org.apache.spark.api.java.function.PairFunction;
import org.geotools.referencing.CRS;
import org.locationtech.jts.geom.Geometry;
import org.opengis.referencing.crs.CoordinateReferenceSystem;
import org.opengis.referencing.operation.MathTransform;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import scala.Tuple2;

import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import static edu.zju.gis.dbfg.model.base.Constant.DEFAULT_CRS;
import static edu.zju.gis.dbfg.model.base.Constant.PYRAMID_CONFIG_4326;

public class Clip  {
    private final static Logger logger = LoggerFactory.getLogger(Clip.class);

    public static void main(String[] args) throws Exception {
        // Setup environment
        logger.info("Setup Spark Context");
        SparkConf conf = new SparkConf();
//        conf.setMaster("local[4]");
//        conf.setAppName("bob");
        JavaSparkContext jsc = new JavaSparkContext(conf);

        //Setup model arguments
        logger.info("Setup arguments");
        ClipArgs mArgs = ClipArgs.initArgs(args);
        if(mArgs == null){
            String msg = "init args failed, exit";
            logger.error(msg);
            throw new ModelFailedException(Clip.class, "main()", msg, args);
        }

        //Transform CRS
        CoordinateReferenceSystem crs1 = CRS.parseWKT(CrsPick.pickByName(mArgs.getCrs1()));
        CoordinateReferenceSystem crs2 = CRS.parseWKT(CrsPick.pickByName(mArgs.getCrs2()));
        CoordinateReferenceSystem targetCRS = DEFAULT_CRS();
        PyramidConfig targetPyramidConfig = PYRAMID_CONFIG_4326();

        MathTransform transform1 = null;
        MathTransform transform2 = null;
        if (!crs1.equals(targetCRS)) {
            transform1 = CRS.findMathTransform(crs1, targetCRS);
        }
        if (!crs2.equals(targetCRS)) {
            transform2 = CRS.findMathTransform(crs2, targetCRS);
        }

        // Read data from datasource and map it to geoItem
        String input1 = mArgs.getInput1();
        String type1 = mArgs.getType1();
        Map<String, Object> attrs1 = new HashMap<>();
        JavaRDD<GeoItem> gi1 = DataReaderRDD.readToRDD(type1, jsc, input1, null, attrs1, transform1);
        JavaPairRDD<String, GeoItem> gi1Pair = DataReaderRDD.mapToPair(gi1, targetPyramidConfig);

        String input2 = mArgs.getInput2();
        String type2 = mArgs.getType2();
        String areaIndex = mArgs.getAreaIndex();
        JavaPairRDD<String, GeoItem> gi2Pair = null;

        if (type2.equals("ES")) {
            gi1Pair.cache();
            List<TileID> tileIDs = gi1Pair.map(t -> TileID.fromString(t._1)).distinct().collect();
            gi2Pair = DataReaderRDD.readToPairRDD(type2, jsc, input2, null, tileIDs, targetPyramidConfig);
        } else {
            JavaRDD<GeoItem> gi2 = DataReaderRDD.readToRDD(type2, jsc, input2, null, null, transform2);
            gi2Pair = DataReaderRDD.mapToPair(gi2, targetPyramidConfig);
        }

        //Clip
        JavaPairRDD<String, Tuple2<GeoItem, GeoItem>> join = gi1Pair.join(gi2Pair);
        JavaPairRDD<String, GeoItem> temp = join.mapToPair(new PairFunction<Tuple2<String, Tuple2<GeoItem, GeoItem>>, String, GeoItem>() {
            @Override
            public Tuple2<String, GeoItem> call(Tuple2<String, Tuple2<GeoItem, GeoItem>> input) throws Exception {
                GeoItem g1 = input._2._1();
                GeoItem g2 = input._2._2();
                Geometry geo1 = g1.getGeometry();
                Geometry geo2 = g2.getGeometry();
                if (!geo1.intersects(geo2)) {
                    return new Tuple2<>("EMPTY", null);
                }
                Geometry intersection = geo2.intersection(geo1);
                Double oarea = Double.valueOf(String.valueOf(g2.getProp().get(areaIndex)));
                GeoItem result = new GeoItem();
                String id = g2.getFid();
                result.setFid(id);
                Map<String, Object> props = new HashMap<>();
                props = g2.getProp();
                //重算面积 采用比例法
                props.put(areaIndex, oarea * (intersection.getArea()/geo2.getArea()));
                result.setProp(props);
                result.setGeometry(intersection);
                return new Tuple2<>(id, result);
            }
        }).filter(new Function<Tuple2<String, GeoItem>, Boolean>() {
            @Override
            public Boolean call(Tuple2<String, GeoItem> input) throws Exception {
                return !input._1.equals("EMPTY");
            }
        });

        temp.cache();
        String output = mArgs.getOutput();
        writeToFile(temp, output, true);

    }

    public static void writeToFile(JavaPairRDD<String, GeoItem> rdd, String filename, Boolean saveAttribute) throws IOException {
        List<Tuple2<String, GeoItem>> featureResult = rdd.collect();
        File featureFile = new File(filename);
        if (featureFile.exists()) featureFile.delete();
        FileWriter featureWriter = new FileWriter(featureFile);
        int fid = 0;
        for (Tuple2<String, GeoItem> feature : featureResult) {
            GeoItem gi = feature._2;
            StringBuilder attribute = new StringBuilder();
            if(saveAttribute)
            {
                Map<String, Object> props = gi.getProp();
                for(String key: props.keySet()){
                    attribute.append(props.get(key));
                    attribute.append("\t");
                }
            }
            Geometry geometry = gi.getGeometry();
            String wkt = geometry.toText();
            fid++;
            featureWriter.write(String.format("%s\t%s%s\n", fid, attribute.toString(), wkt));
        }
        featureWriter.flush();
        featureWriter.close();
    }

}
